//===----------------------------------------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2024 Apple Inc. and the Swift.org project authors
// Licensed under Apache License v2.0
//
// See LICENSE.txt for license information
// See CONTRIBUTORS.txt for the list of Swift.org project authors
//
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//

import SwiftSyntax
import SwiftSyntaxBuilder

extension FFMSwift2JavaGenerator {
  package func writeSwiftThunkSources() throws {
    var printer = CodePrinter()
    try writeSwiftThunkSources(printer: &printer)
  }

  package func writeSwiftExpectedEmptySources() throws {
    let pendingFileCount = self.expectedOutputSwiftFileNames.count
    guard pendingFileCount > 0 else {
      return // no need to write any empty files, yay
    }
    
    log.info("[swift-java] Write empty [\(self.expectedOutputSwiftFileNames.count)] 'expected' files in: \(swiftOutputDirectory)/")
    
    for expectedFileName in self.expectedOutputSwiftFileNames {
      log.info("Write SwiftPM-'expected' empty file: \(expectedFileName.bold)")

      var printer = CodePrinter()
      printer.print("// Empty file generated on purpose")
      _ = try printer.writeContents(
        outputDirectory: self.swiftOutputDirectory,
        javaPackagePath: nil,
        filename: expectedFileName)
    }
  }

  package func writeSwiftThunkSources(printer: inout CodePrinter) throws {
    let moduleFilenameBase = "\(self.swiftModuleName)Module+SwiftJava"
    let moduleFilename = "\(moduleFilenameBase).swift"
    do {
      log.debug("Printing contents: \(moduleFilename)")

      try printGlobalSwiftThunkSources(&printer)

      if let outputFile = try printer.writeContents(
        outputDirectory: self.swiftOutputDirectory,
        javaPackagePath: nil,
        filename: moduleFilename) {
        log.info("Generated: \(moduleFilenameBase.bold).swift (at \(outputFile.absoluteString))")
        self.expectedOutputSwiftFileNames.remove(moduleFilename)
      }
    } catch {
      log.warning("Failed to write to Swift thunks: \(moduleFilename)")
    }

    // === All types
    // We have to write all types to their corresponding output file that matches the file they were declared in,
    // because otherwise SwiftPM plugins will not pick up files apropriately -- we expect 1 output +SwiftJava.swift file for every input.
    for group: (key: String, value: [Dictionary<String, ImportedNominalType>.Element]) in Dictionary(grouping: self.analysis.importedTypes, by: { $0.value.sourceFilePath }) {
      log.warning("Writing types in file group: \(group.key): \(group.value.map(\.key))")

      let importedTypesForThisFile = group.value
        .map(\.value)
        .sorted(by: { $0.qualifiedName < $1.qualifiedName })

      let inputFileName = "\(group.key)".split(separator: "/").last ?? "__Unknown.swift"
      let filename = "\(inputFileName)".replacing(".swift", with: "+SwiftJava.swift")

      for ty in importedTypesForThisFile {
        log.info("Printing Swift thunks for type: \(ty.qualifiedName.bold)")
        printer.printSeparator("Thunks for \(ty.qualifiedName)")

        do {
          try printSwiftThunkSources(&printer, ty: ty)
        } catch {
          log.warning("Failed to print to Swift thunks for type'\(ty.qualifiedName)' to '\(filename)', error: \(error)")
        }
        
      }

      log.warning("Write Swift thunks file: \(filename.bold)")
      do {
        if let outputFile = try printer.writeContents(
          outputDirectory: self.swiftOutputDirectory,
          javaPackagePath: nil,
          filename: filename) {
          log.info("Done writing Swift thunks to: \(outputFile.absoluteString)")
          self.expectedOutputSwiftFileNames.remove(filename)
        }
      } catch {
        log.warning("Failed to write to Swift thunks: \(filename), error: \(error)")
      }
    }
  }

  public func printGlobalSwiftThunkSources(_ printer: inout CodePrinter) throws {
    let stt = SwiftThunkTranslator(self)

    printer.print(
      """
      // Generated by swift-java

      import SwiftRuntimeFunctions

      """)

    printSwiftThunkImports(&printer)

    for thunk in stt.renderGlobalThunks() {
      printer.print(thunk)
      printer.println()
    }
  }

  public func printSwiftThunkSources(_ printer: inout CodePrinter, decl: ImportedFunc) {
    let stt = SwiftThunkTranslator(self)

    for thunk in stt.render(forFunc: decl) {
      printer.print(thunk)
      printer.println()
    }
  }

  package func printSwiftThunkSources(_ printer: inout CodePrinter, ty: ImportedNominalType) throws {
    let stt = SwiftThunkTranslator(self)

    printer.print(
      """
      // Generated by swift-java

      import SwiftRuntimeFunctions

      """
    )

    printSwiftThunkImports(&printer)

    for thunk in stt.renderThunks(forType: ty) {
      printer.print("\(thunk)")
      printer.print("")
    }
  }

  func printSwiftThunkImports(_ printer: inout CodePrinter) {
    let mainSymbolSourceModules = Set(
      self.lookupContext.symbolTable.importedModules.values.filter { $0.alternativeModules?.isMainSourceOfSymbols ?? false }.map(\.moduleName)
    )
    
    for module in self.lookupContext.symbolTable.importedModules.keys.sorted() {
      guard module != "Swift" else {
        continue
      }

      guard let alternativeModules = self.lookupContext.symbolTable.importedModules[module]?.alternativeModules else {
        printer.print("import \(module)")
        continue
      }

      // Try to print only on main module from relation chain as it has every other module.
      guard !mainSymbolSourceModules.isDisjoint(with: alternativeModules.moduleNames) || alternativeModules.isMainSourceOfSymbols else {
        if !alternativeModules.isMainSourceOfSymbols {
          printer.print("import \(module)")
        }
        continue
      }
      
      var importGroups: [String: [String]] = [:]
      for name in alternativeModules.moduleNames {
        guard let otherModule = self.lookupContext.symbolTable.importedModules[name] else { continue }

        let groupKey = otherModule.requiredAvailablityOfModuleWithName ?? otherModule.moduleName
        importGroups[groupKey, default: []].append(otherModule.moduleName)
      }

      for (index, group) in importGroups.keys.sorted().enumerated() {
        if index > 0  && importGroups.keys.count > 1 {
          printer.print("#elseif canImport(\(group))")
        } else {
          printer.print("#if canImport(\(group))")
        }
        
        for groupModule in importGroups[group] ?? [] {
          printer.print("import \(groupModule)")
        }
      }

      if (importGroups.keys.isEmpty) {
        printer.print("import \(module)")
      } else {
        printer.print("#else")
        printer.print("import \(module)")
        printer.print("#endif")
      }
    }
    printer.println()
  }
}

struct SwiftThunkTranslator {

  let st: FFMSwift2JavaGenerator

  init(_ st: FFMSwift2JavaGenerator) {
    self.st = st
  }

  func renderGlobalThunks() -> [DeclSyntax] {
    var decls: [DeclSyntax] = []
    decls.reserveCapacity(
      st.analysis.importedGlobalVariables.count + st.analysis.importedGlobalFuncs.count
    )

    for decl in st.analysis.importedGlobalVariables {
      decls.append(contentsOf: render(forFunc: decl))
    }

    for decl in st.analysis.importedGlobalFuncs {
      decls.append(contentsOf: render(forFunc: decl))
    }

    return decls
  }

  /// Render all the thunks that make Swift methods accessible to Java.
  func renderThunks(forType nominal: ImportedNominalType) -> [DeclSyntax] {
    var decls: [DeclSyntax] = []
    decls.reserveCapacity(
      1 + nominal.initializers.count + nominal.variables.count + nominal.methods.count
    )

    decls.append(renderSwiftTypeAccessor(nominal))

    for decl in nominal.initializers {
      decls.append(contentsOf: render(forFunc: decl))
    }

    for decl in nominal.variables {
      decls.append(contentsOf: render(forFunc: decl))
    }

    for decl in nominal.methods {
      decls.append(contentsOf: render(forFunc: decl))
    }

    return decls
  }

  /// Accessor to get the `T.self` of the Swift type, without having to rely on mangled name lookups.
  func renderSwiftTypeAccessor(_ nominal: ImportedNominalType) -> DeclSyntax {
    let funcName = SwiftKitPrinting.Names.getType(
      module: st.swiftModuleName,
      nominal: nominal)

    return
      """
      @_cdecl("\(raw: funcName)")
      public func \(raw: funcName)() -> UnsafeMutableRawPointer /* Any.Type */ {
        return unsafeBitCast(\(raw: nominal.swiftNominal.qualifiedName).self, to: UnsafeMutableRawPointer.self)
      }
      """
  }

  func render(forFunc decl: ImportedFunc) -> [DeclSyntax] {
    st.log.trace("Rendering thunks for: \(decl.displayName)")

    let thunkName = st.thunkNameRegistry.functionThunkName(decl: decl)
    guard let translated = st.translatedDecl(for: decl) else {
      return []
    }

    let thunkFunc = translated.loweredSignature.cdeclThunk(
      cName: thunkName,
      swiftAPIName: decl.name,
      as: decl.apiKind
    )
    return [DeclSyntax(thunkFunc)]
  }
}
